package utils

import (
	"context"
	"fmt"
	"sync"
	"time"

)

const (
	gDefaultWorkerNum    = 8
	gDefaultTaskCapacity = 1024
	gAddTimeoutSecond    = 5
	gExecTimeoutSecond   = 120 // 120s
)

// task function.
type Task func() interface{}

// TaskResult
type TaskResult struct {
	Result interface{}
	Error  error
}

// TaskWithResult
type TaskWithResult struct {
	Task     Task
	ResultCh chan TaskResult
	Timeout  time.Duration // 0 is unlimited.
}

// goroutine pool interface.
type IGoroutinePool interface {
	// Start goroutine pool's task.
	Start()

	// Add task with result.
	AddTaskWithResult(task TaskWithResult) error

	// Add task without result and no timeout.
	AddTask(task Task) error

	// Close and wait to close.
	Wait()
}

// goroutine pool.
type goroutinePool struct {
	tasks      chan TaskWithResult
	numWorkers int
	wg         sync.WaitGroup
}

// NewGoroutinePool create goroutine pool.
func NewGoroutinePool(numWorkers int, taskCapacity int) IGoroutinePool {
	if numWorkers <= 0 {
		numWorkers = gDefaultWorkerNum
	}
	if taskCapacity <= 0 {
		taskCapacity = gDefaultTaskCapacity
	}

	pool := &goroutinePool{
		tasks:      make(chan TaskWithResult, taskCapacity),
		numWorkers: numWorkers,
	}
	return pool
}

// Start start to run.
func (p *goroutinePool) Start() {
	for i := 0; i < p.numWorkers; i++ {
		p.wg.Add(1)
		go p.worker()
	}
}

// AddTaskWithResult AddTask add task with result.
func (p *goroutinePool) AddTaskWithResult(task TaskWithResult) error {
	if task.Task == nil {
		return fmt.Errorf("task is nil")
	}

	if task.Timeout < 0 {
		task.Timeout = gExecTimeoutSecond * time.Second
	}

	select {
	case p.tasks <- task:
		return nil
	case <-time.After(time.Second * gAddTimeoutSecond):
		return fmt.Errorf("add time out")
	}
}

// AddTask adds task without result.
func (p *goroutinePool) AddTask(task Task) error {
	if task == nil {
		return fmt.Errorf("task is nil")
	}

	var taskWithResult TaskWithResult
	taskWithResult.ResultCh = nil
	taskWithResult.Task = task
	taskWithResult.Timeout = 0

	select {
	case p.tasks <- taskWithResult:
		return nil
	case <-time.After(time.Second * gAddTimeoutSecond):
		return fmt.Errorf("add time out")
	}
}

// Wait wait close and wait.
func (p *goroutinePool) Wait() {
	close(p.tasks)
	p.wg.Wait()
}

// worker function.
func (p *goroutinePool) worker() {
	defer p.wg.Done()
	for task := range p.tasks {
		func(task TaskWithResult) {
			defer func() {
				if r := recover(); r != nil {
					if task.ResultCh != nil {
						task.ResultCh <- TaskResult{
							Result: nil,
							Error:  fmt.Errorf("recovered from panic in task: %v", r),
						}
					}
				}
			}()

			var ctx context.Context
			var cancel context.CancelFunc

			// Create a context with the specified timeout if timeout isn't zero
			if task.Timeout == 0 {
				ctx, cancel = context.WithCancel(context.Background())
			} else {
				ctx, cancel = context.WithTimeout(context.Background(), task.Timeout)
			}

			defer cancel()

			var result interface{}
			var err error
			doneCh := make(chan struct{})

			go func() {
				defer close(doneCh)
				result = task.Task()
			}()

			if task.Timeout == 0 {
				<-doneCh
				err = nil
			} else {
				select {
				case <-doneCh:
					err = nil
				case <-ctx.Done():
					result = nil
					err = fmt.Errorf("task execution timeout")
				}
			}

			if task.ResultCh != nil {
				task.ResultCh <- TaskResult{Result: result, Error: err}
			}
		}(task)
	}
}
